import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.tri as tri


# 自定义的不规则函数
def irregular_function(x, y):
    return x * y * np.sin(x) * (x / 5 + np.cos(y))


# 创建彩色的应力分布云图
def plot_colorful_stress_distribution(ax):
    # 生成一些示例数据，你需要替换这部分数据为你的实际数据
    x = np.linspace(0, 10, 100)
    y = np.linspace(0, 10, 100)
    X, Y = np.meshgrid(x, y)
    Z = irregular_function(X, Y)

    # 绘制彩色的应力分布云图
    surf = ax.plot_surface(X, Y, Z, cmap="viridis")
    plt.colorbar(surf, ax=ax, label="Stress")
    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    ax.set_title("Stress Distribution")


# 创建在彩色应力分布云图基础上均匀分布黑色节点的图
def plot_uniform_black_nodes(ax):
    # 使用相同的示例数据生成黑色节点
    x = np.linspace(0, 10, 100)
    y = np.linspace(0, 10, 100)
    X, Y = np.meshgrid(x, y)

    # 随机生成一些均匀分布的黑色节点的坐标
    num_nodes = 100
    x_nodes = np.random.uniform(0, 10, num_nodes)
    y_nodes = np.random.uniform(0, 10, num_nodes)

    # 绘制彩色的应力分布云图
    surf = ax.plot_wireframe(X, Y, irregular_function(X, Y), color="black", linewidth=0.5)

    # 绘制黑色节点
    ax.scatter(x_nodes, y_nodes, irregular_function(x_nodes, y_nodes), color="red", s=20, label="Nodes")

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    ax.set_title("Nodes")


# 创建一个水平排列的图形
fig = plt.figure(figsize=(16, 6))

ax1 = fig.add_subplot(1, 2, 1, projection="3d")  # 1行2列，第1个图
plot_colorful_stress_distribution(ax1)

ax2 = fig.add_subplot(1, 2, 2, projection="3d")  # 1行2列，第2个图
plot_uniform_black_nodes(ax2)

plt.tight_layout()  # 调整布局使得子图之间不重叠
plt.show()
